"""
    Using Pytorch to simulate render pipline in texture-space.
"""

import torch
import torch.nn as nn
import torch.nn.functional as F
from pytorch_util import gaussian_kernel_generation
import pytorch_config as C
import math


class TextureSpaceRender(nn.Module):
    def __init__(self,
                 light_direction: torch.Tensor,
                 light_radiance: torch.Tensor,
                 view_direction: torch.Tensor,
                 ambient_illumination: float = 0.2,
                 specular_exponent: float = 16,
                 exposure: float = 1,
                 gamma: float = 2.2,
                 using_specular_normal_when_diffuse: bool = True,
                 using_sss: bool = True,
                 using_gamma_correction: bool = False,
                 using_hdr: bool = True):
        """
            Initialize renderer.
            Perhaps convolution kernel need to be initialized here.
        """
        super().__init__()
        self.light_d = F.normalize(light_direction, p=2, dim=0).to(torch.double)
        self.light_r = light_radiance.to(torch.double)
        self.ambient_i = ambient_illumination
        self.view_d = F.normalize(view_direction, p=2, dim=0).to(torch.double)
        self.s_e = specular_exponent
        self.exposure = exposure
        self.gamma = gamma
        self.using_specular_normal_when_diffuse = using_specular_normal_when_diffuse
        self.using_sss = using_sss
        self.using_gamma_correction = using_gamma_correction
        self.using_hdr = using_hdr

        '''
            Initialize gaussian blur kernels
        '''
        self.gaussian_kernel_weight2 = gaussian_kernel_generation(
            kernel_size=C.gaussian_kernel_size2,
            sigma=math.sqrt(C.gaussian_variance2)).to(torch.double)
        self.blur_weight2 = torch.Tensor(C.gaussian_blur_weight2).to(torch.double)
        self.gaussian_kernel_weight3 = gaussian_kernel_generation(
            kernel_size=C.gaussian_kernel_size3,
            sigma=math.sqrt(C.gaussian_variance3)).to(torch.double)
        self.blur_weight3 = torch.Tensor(C.gaussian_blur_weight3).to(torch.double)
        self.gaussian_kernel_weight4 = gaussian_kernel_generation(
            kernel_size=C.gaussian_kernel_size4,
            sigma=math.sqrt(C.gaussian_variance4)).to(torch.double)
        self.blur_weight4 = torch.Tensor(C.gaussian_blur_weight4).to(torch.double)
        self.gaussian_kernel_weight5 = gaussian_kernel_generation(
            kernel_size=C.gaussian_kernel_size5,
            sigma=math.sqrt(C.gaussian_variance5)).to(torch.double)
        self.blur_weight5 = torch.Tensor(C.gaussian_blur_weight5).to(torch.double)
        self.gaussian_kernel_weight6 = gaussian_kernel_generation(
            kernel_size=C.gaussian_kernel_size6,
            sigma=math.sqrt(C.gaussian_variance6)).to(torch.double)
        self.blur_weight6 = torch.Tensor(C.gaussian_blur_weight6).to(torch.double)

    def forward(self,
                diffuse_albedo: torch.Tensor,
                diffuse_normal: torch.Tensor,
                specular_albedo: torch.Tensor,
                specular_normal: torch.Tensor) -> torch.Tensor:
        """
            Map should be able to use straightly, instead of preprocessing there.
            map format : (N, H, W, C)
                N : batch size
                W : image width
                H : image height
                C : image channels
        """
        '''
            Calculate diffuse component.
        '''
        # F.relu has the same effect as max(..., 0), but differentiable.
        if self.using_specular_normal_when_diffuse:
            diffuse_intensity = F.relu((self.light_d * specular_normal).sum(dim=-1))
        else:
            diffuse_intensity = F.relu((self.light_d * diffuse_normal).sum(dim=-1))
        diff_factor = torch.unsqueeze(diffuse_intensity, dim=-1).expand(-1, -1, -1, 3) * self.ambient_i * self.light_r
        diffuse_component = diff_factor * diffuse_albedo

        '''
            If enable subsurface scattering, apply following process.
        '''
        if self.using_sss:
            """
                Convert diffuse_component into format as (N * C * H * W).
            """
            reshaped_diffuse_component = diffuse_component.permute(0, 3, 1, 2)

            blur2 = F.conv2d(reshaped_diffuse_component,
                             self.gaussian_kernel_weight2,
                             stride=1,
                             padding=C.gaussian_kernel_size2 // 2,
                             groups=3).permute(0, 2, 3, 1)
            blur3 = F.conv2d(reshaped_diffuse_component,
                             self.gaussian_kernel_weight3,
                             stride=1,
                             padding=C.gaussian_kernel_size3 // 2,
                             groups=3).permute(0, 2, 3, 1)
            blur4 = F.conv2d(reshaped_diffuse_component,
                             self.gaussian_kernel_weight4,
                             stride=1,
                             padding=C.gaussian_kernel_size4 // 2,
                             groups=3).permute(0, 2, 3, 1)
            blur5 = F.conv2d(reshaped_diffuse_component,
                             self.gaussian_kernel_weight5,
                             stride=1,
                             padding=C.gaussian_kernel_size5 // 2,
                             groups=3).permute(0, 2, 3, 1)
            blur6 = F.conv2d(reshaped_diffuse_component,
                             self.gaussian_kernel_weight6,
                             stride=1,
                             padding=C.gaussian_kernel_size6 // 2,
                             groups=3).permute(0, 2, 3, 1)
            reshaped_diffuse_component = reshaped_diffuse_component.permute(0, 2, 3, 1)
            blur_diffuse_component = reshaped_diffuse_component * torch.Tensor(C.gaussian_blur_weight1)
            blur_diffuse_component += blur2 * torch.Tensor(C.gaussian_blur_weight2)
            blur_diffuse_component += blur3 * torch.Tensor(C.gaussian_blur_weight3)
            blur_diffuse_component += blur4 * torch.Tensor(C.gaussian_blur_weight4)
            blur_diffuse_component += blur5 * torch.Tensor(C.gaussian_blur_weight5)
            blur_diffuse_component += blur6 * torch.Tensor(C.gaussian_blur_weight6)
            diffuse_component = blur_diffuse_component

        '''
            Calculate specular component
        '''
        half = F.normalize(self.light_d + self.view_d, 2, dim=0)
        specular_intensity = F.relu((specular_normal * half).sum(dim=-1))
        specular_factor = torch.unsqueeze(specular_intensity, dim=-1).expand(-1, -1, -1, 3).pow(self.s_e) * self.light_r
        specular_component = specular_factor * specular_albedo

        '''
            Map HDR color into [0, 1]
        '''
        if self.using_hdr:
            hdr_color = diffuse_component + specular_component * 0.3
            mapped_color = 1 - torch.exp(-hdr_color * self.exposure)
        else:
            mapped_color = diffuse_component * 0.7 + specular_component * 0.3

        '''
            If gamma correction is needed, apply it.
        '''
        if self.using_gamma_correction:
            mapped_color = mapped_color ** (1 / self.gamma)

        return mapped_color


if __name__ == "__main__":
    ...
